// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <sstream>

#include <nebula/compute/api_scalar.h>
#include <nebula/compute/kernels/common_internal.h>

namespace nebula::compute::internal {

// ----------------------------------------------------------------------
// String transformation base classes

    constexpr int64_t kStringTransformError = -1;

    struct StringTransformBase {
        virtual ~StringTransformBase() = default;

        virtual turbo::Status PreExec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            return turbo::OkStatus();
        }

        // Return the maximum total size of the output in codeunits (i.e. bytes)
        // given input characteristics.
        virtual int64_t MaxCodeunits(int64_t ninputs, int64_t input_ncodeunits) {
            return input_ncodeunits;
        }

        virtual turbo::Status InvalidInputSequence() {
            return turbo::invalid_argument_error("Invalid UTF8 sequence in input");
        }
    };

    template<typename offset_type>
    static int64_t GetVarBinaryValuesLength(const ArraySpan &span) {
        const offset_type *offsets = span.get_values<offset_type>(1);
        return span.length > 0 ? offsets[span.length] - offsets[0] : 0;
    }

/// Kernel exec generator for unary string transforms. Types of template
/// parameter StringTransform need to define a transform method with the
/// following signature:
///
/// int64_t Transform(const uint8_t* input, int64_t input_string_ncodeunits,
///                   uint8_t* output);
///
/// where
///   * `input` is the input sequence (binary or string)
///   * `input_string_ncodeunits` is the length of input sequence in codeunits
///   * `output` is the output sequence (binary or string)
///
/// and returns the number of codeunits of the `output` sequence or a negative
/// value if an invalid input sequence is detected.
    template<typename Type, typename StringTransform>
    struct StringTransformExecBase {
        using offset_type = typename Type::offset_type;
        using ArrayType = typename TypeTraits<Type>::ArrayType;

        static turbo::Status execute(KernelContext *ctx, StringTransform *transform,
                                     const ExecSpan &batch, ExecResult *out) {
            const ArraySpan &input = batch[0].array;
            auto offsets = input.get_values<offset_type>(1);
            const uint8_t *input_data = input.buffers[2].data;

            const int64_t input_ncodeunits = GetVarBinaryValuesLength<offset_type>(input);
            const int64_t max_output_ncodeunits =
                    transform->MaxCodeunits(input.length, input_ncodeunits);
            TURBO_RETURN_NOT_OK(check_output_capacity(max_output_ncodeunits));

            ArrayData *output = out->array_data().get();
            TURBO_MOVE_OR_RAISE(auto values_buffer, ctx->allocate(max_output_ncodeunits));
            output->buffers[2] = values_buffer;

            // String offsets are preallocated
            offset_type *output_string_offsets = output->GetMutableValues<offset_type>(1);
            uint8_t *output_str = output->buffers[2]->mutable_data();
            offset_type output_ncodeunits = 0;
            output_string_offsets[0] = output_ncodeunits;
            for (int64_t i = 0; i < input.length; i++) {
                if (!input.is_null(i)) {
                    const uint8_t *input_string = input_data + offsets[i];
                    offset_type input_string_ncodeunits = offsets[i + 1] - offsets[i];
                    auto encoded_nbytes = static_cast<offset_type>(transform->Transform(
                            input_string, input_string_ncodeunits, output_str + output_ncodeunits));
                    if (encoded_nbytes < 0) {
                        return transform->InvalidInputSequence();
                    }
                    output_ncodeunits += encoded_nbytes;
                }
                output_string_offsets[i + 1] = output_ncodeunits;
            }
            DKCHECK_LE(output_ncodeunits, max_output_ncodeunits);

            // Trim the codepoint buffer, since we may have allocated too much
            return values_buffer->resize(output_ncodeunits, /*shrink_to_fit=*/true);
        }

        static turbo::Status check_output_capacity(int64_t ncodeunits) {
            if (ncodeunits > std::numeric_limits<offset_type>::max()) {
                return turbo::resource_exhausted_error(
                        "turbo::Result might not fit in a 32bit utf8 array, convert to large_utf8");
            }
            return turbo::OkStatus();
        }
    };

    template<typename Type, typename StringTransform>
    struct StringTransformExec : public StringTransformExecBase<Type, StringTransform> {
        using StringTransformExecBase<Type, StringTransform>::execute;

        static turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            StringTransform transform;
            TURBO_RETURN_NOT_OK(transform.PreExec(ctx, batch, out));
            return execute(ctx, &transform, batch, out);
        }
    };

    template<typename Type, typename StringTransform>
    struct StringTransformExecWithState
            : public StringTransformExecBase<Type, StringTransform> {
        using State = typename StringTransform::State;
        using StringTransformExecBase<Type, StringTransform>::execute;

        static turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            StringTransform transform(State::Get(ctx));
            TURBO_RETURN_NOT_OK(transform.PreExec(ctx, batch, out));
            return execute(ctx, &transform, batch, out);
        }
    };

    template<template<typename> class ExecFunctor>
    void make_unary_string_batch_kernel(
            std::string name, FunctionRegistry *registry, FunctionDoc doc,
            MemAllocation::type mem_allocation = MemAllocation::PREALLOCATE) {
        auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));
        for (const auto &ty: string_types()) {
            auto exec = generate_var_binary_to_var_binary<ExecFunctor>(ty);
            ScalarKernel kernel{{ty}, ty, std::move(exec)};
            kernel.mem_allocation = mem_allocation;
            KCHECK_OK(func->add_kernel(std::move(kernel)));
        }
        KCHECK_OK(registry->add_function(std::move(func)));
    }

    template<template<typename> class ExecFunctor>
    void make_unary_string_batch_kernel_with_state(
            std::string name, FunctionRegistry *registry, FunctionDoc doc,
            MemAllocation::type mem_allocation = MemAllocation::PREALLOCATE) {
        auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));
        {
            using t32 = ExecFunctor<StringType>;
            ScalarKernel kernel{{utf8()}, utf8(), t32::Exec, t32::State::init};
            kernel.mem_allocation = mem_allocation;
            KCHECK_OK(func->add_kernel(std::move(kernel)));
        }
        {
            using t64 = ExecFunctor<LargeStringType>;
            ScalarKernel kernel{{large_utf8()}, large_utf8(), t64::Exec, t64::State::init};
            kernel.mem_allocation = mem_allocation;
            KCHECK_OK(func->add_kernel(std::move(kernel)));
        }
        KCHECK_OK(registry->add_function(std::move(func)));
    }

    // ----------------------------------------------------------------------
    // Predicates and classification

    // Defined in scalar_string_utf8.cc.
    void EnsureUtf8LookupTablesFilled();

    static FunctionDoc StringPredicateDoc(std::string summary, std::string description) {
        return FunctionDoc{std::move(summary), std::move(description), {"strings"}};
    }

    static inline FunctionDoc StringClassifyDoc(std::string class_summary,
                                                std::string class_desc, bool non_empty) {
        std::string summary, description;
        {
            std::stringstream ss;
            ss << "Classify strings as " << class_summary;
            summary = ss.str();
        }
        {
            std::stringstream ss;
            if (non_empty) {
                ss
                        << ("For each string in `strings`, emit true iff the string is non-empty\n"
                            "and consists only of ");
            } else {
                ss
                        << ("For each string in `strings`, emit true iff the string consists only\n"
                            "of ");
            }
            ss << class_desc << ".  Null strings emit null.";
            description = ss.str();
        }
        return StringPredicateDoc(std::move(summary), std::move(description));
    }

    template<typename Type, typename Predicate>
    struct StringPredicateFunctor {
        static turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            turbo::Status st = turbo::OkStatus();
            EnsureUtf8LookupTablesFilled();
            const ArraySpan &input = batch[0].array;
            ArrayIterator<Type> input_it(input);
            ArraySpan *out_arr = out->array_span_mutable();
            ::nebula::internal::GenerateBitsUnrolled(
                    out_arr->buffers[1].data, out_arr->offset, input.length, [&]() -> bool {
                        std::string_view val = input_it();
                        return Predicate::Call(ctx, reinterpret_cast<const uint8_t *>(val.data()),
                                               val.size(), &st);
                    });
            return st;
        }
    };

    template<typename Predicate>
    void AddUnaryStringPredicate(std::string name, FunctionRegistry *registry,
                                 FunctionDoc doc) {
        auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));
        for (const auto &ty: string_types()) {
            auto exec = generate_var_binary_to_var_binary<StringPredicateFunctor, Predicate>(ty);
            KCHECK_OK(func->add_kernel({ty}, boolean(), std::move(exec)));
        }
        KCHECK_OK(registry->add_function(std::move(func)));
    }

    // ----------------------------------------------------------------------
    // Slicing

    struct StringSliceTransformBase : public StringTransformBase {
        using State = OptionsWrapper<SliceOptions>;

        const SliceOptions *options;

        StringSliceTransformBase() = default;

        explicit StringSliceTransformBase(const SliceOptions &options) : options{&options} {}

        turbo::Status PreExec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) override {
            options = &State::Get(ctx);
            if (options->step == 0) {
                return turbo::invalid_argument_error("slice step cannot be zero");
            }
            return turbo::OkStatus();
        }
    };

    struct ReplaceStringSliceTransformBase : public StringTransformBase {
        using State = OptionsWrapper<ReplaceSliceOptions>;

        const ReplaceSliceOptions *options;

        explicit ReplaceStringSliceTransformBase(const ReplaceSliceOptions &options)
                : options{&options} {}

        int64_t MaxCodeunits(int64_t ninputs, int64_t input_ncodeunits) override {
            return ninputs * options->replacement.size() + input_ncodeunits;
        }
    };

// ----------------------------------------------------------------------
// Splitting

    template<typename Options>
    struct StringSplitFinderBase {
        virtual ~StringSplitFinderBase() = default;

        virtual turbo::Status PreExec(const Options &options) { return turbo::OkStatus(); }

        // Derived classes should also define these methods:
        //   static bool Find(const uint8_t* begin, const uint8_t* end,
        //                    const uint8_t** separator_begin,
        //                    const uint8_t** separator_end,
        //                    const SplitPatternOptions& options);
        //
        //   static bool FindReverse(const uint8_t* begin, const uint8_t* end,
        //                           const uint8_t** separator_begin,
        //                           const uint8_t** separator_end,
        //                           const SplitPatternOptions& options);
    };

    template<typename Type, typename ListType, typename SplitFinder,
            typename Options = typename SplitFinder::Options>
    struct StringSplitExec {
        using string_offset_type = typename Type::offset_type;
        using list_offset_type = typename ListType::offset_type;
        using ArrayType = typename TypeTraits<Type>::ArrayType;
        using ArrayListType = typename TypeTraits<ListType>::ArrayType;
        using ListScalarType = typename TypeTraits<ListType>::ScalarType;
        using ScalarType = typename TypeTraits<Type>::ScalarType;
        using BuilderType = typename TypeTraits<Type>::BuilderType;
        using ListOffsetsBuilderType = TypedBufferBuilder<list_offset_type>;
        using State = OptionsWrapper<Options>;

        // Keep the temporary storage across individual values, to minimize reallocations
        std::vector<std::string_view> parts;
        Options options;

        explicit StringSplitExec(const Options &options) : options(options) {}

        static turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            return StringSplitExec{State::Get(ctx)}.execute(ctx, batch, out);
        }

        turbo::Status execute(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            SplitFinder finder;
            TURBO_RETURN_NOT_OK(finder.PreExec(options));
            // TODO(wesm): refactor to not require creating ArrayData
            const ArrayType input(batch[0].array.to_array_data());

            BuilderType builder(input.type(), ctx->memory_pool());
            // A slight overestimate of the data needed
            TURBO_RETURN_NOT_OK(builder.reserve_data(input.total_values_length()));
            // The minimum amount of strings needed
            TURBO_RETURN_NOT_OK(builder.resize(input.length() - input.null_count()));

            ArrayData *output_list = out->array_data().get();
            // List offsets were preallocated
            auto *list_offsets = output_list->GetMutableValues<list_offset_type>(1);
            DKCHECK_NE(list_offsets, nullptr);
            // Initial value
            *list_offsets++ = 0;
            for (int64_t i = 0; i < input.length(); ++i) {
                if (!input.is_null(i)) {
                    TURBO_RETURN_NOT_OK(SplitString(input.get_view(i), &finder, &builder));
                    if (TURBO_UNLIKELY(builder.length() >
                                       std::numeric_limits<list_offset_type>::max())) {
                        return turbo::resource_exhausted_error("List offset does not fit into 32 bit");
                    }
                }
                *list_offsets++ = static_cast<list_offset_type>(builder.length());
            }
            // Assign string array to list child data
            std::shared_ptr<Array> string_array;
            TURBO_RETURN_NOT_OK(builder.finish(&string_array));
            output_list->child_data.push_back(string_array->data());
            return turbo::OkStatus();
        }

        turbo::Status SplitString(std::string_view s, SplitFinder *finder, BuilderType *builder) {
            const uint8_t *begin = reinterpret_cast<const uint8_t *>(s.data());
            const uint8_t *end = begin + s.length();

            int64_t max_splits = options.max_splits;
            // if there is no max splits, reversing does not make sense (and is probably less
            // efficient), but is useful for testing
            if (options.reverse) {
                // note that i points 1 further than the 'current'
                const uint8_t *i = end;
                // we will record the parts in reverse order
                parts.clear();
                if (max_splits > -1) {
                    parts.reserve(max_splits + 1);
                }
                while (max_splits != 0) {
                    const uint8_t *separator_begin, *separator_end;
                    // find with whatever algo the part we will 'cut out'
                    if (finder->FindReverse(begin, i, &separator_begin, &separator_end, options)) {
                        parts.emplace_back(reinterpret_cast<const char *>(separator_end),
                                           i - separator_end);
                        i = separator_begin;
                        max_splits--;
                    } else {
                        // if we cannot find a separator, we're done
                        break;
                    }
                }
                parts.emplace_back(reinterpret_cast<const char *>(begin), i - begin);
                // now we do the copying
                for (auto it = parts.rbegin(); it != parts.rend(); ++it) {
                    TURBO_RETURN_NOT_OK(builder->append(*it));
                }
            } else {
                const uint8_t *i = begin;
                while (max_splits != 0) {
                    const uint8_t *separator_begin, *separator_end;
                    // find with whatever algo the part we will 'cut out'
                    if (finder->Find(i, end, &separator_begin, &separator_end, options)) {
                        // the part till the beginning of the 'cut'
                        TURBO_RETURN_NOT_OK(
                                builder->append(i, static_cast<string_offset_type>(separator_begin - i)));
                        i = separator_end;
                        max_splits--;
                    } else {
                        // if we cannot find a separator, we're done
                        break;
                    }
                }
                // trailing part
                TURBO_RETURN_NOT_OK(builder->append(i, static_cast<string_offset_type>(end - i)));
            }
            return turbo::OkStatus();
        }
    };

    using StringSplitState = OptionsWrapper<SplitOptions>;

}  // namespace nebula::compute::internal
